# -- coding: utf-8 --
""" 动量方程求解

"""

# third party lib
import numpy as np

from utils.constant import Physics, InterpMode
from utils.numeric import grid_interp, penta_diagonal_solve, grid_extrap, difference

def cal_moment_L(params,
                 liquid_wall_friction_force: np.ndarray, 

                 gas_liquid_friction_force: np.ndarray, 
                 liquid_vapour_mass_transfer_rate: np.ndarray, 
                 V_i_segment: np.ndarray, 
                 transition_velocity_from_liquid_film: np.ndarray, 
                 liquid_to_droplet_mass_transfer_rate: np.ndarray, 
                 droplet_to_liquid_mass_transfer_rate: np.ndarray, 
                 droplet_velocity: np.ndarray) -> np.ndarray:
    """ 动量方程计算,液相

    Args:
        params: dict 参数字典
        liquid_wall_friction_force: np.ndarray 液相壁面摩擦力
        gas_liquid_friction_force: np.ndarray 气液摩擦力
        liquid_vapour_mass_transfer_rate: np.ndarray 液相蒸气质量传递速率
        V_i_segment: np.ndarray 液相到液滴质量传递速率
        transition_velocity_from_liquid_film: np.ndarray 来自液膜蒸发的相变速度
        liquid_to_droplet_mass_transfer_rate: np.ndarray 液相到液滴质量传递速率
        droplet_to_liquid_mass_transfer_rate: np.ndarray 液滴到液相质量传递速率
        droplet_velocity: np.ndarray 液滴速度

    Returns:
        V_L_segment_update: np.ndarray:, 液相速度
    """
    dx = params["dx"]
    dt = params["dt"]
    p_node = params["p_node"]
    cross_section_area_segment = params["cross_section_area_segment"]
    tilt_angle_segment = params["tilt_angle_segment"]
    D_segment = params["D_segment"]
    # 速度相关变量
    V_L_segment = params["V_L_segment"]
    V_g_segment = params["V_g_segment"]
    
    # 密度相关变量
    rho_L_node = params["rho_L_node"]
    rho_L_segment = grid_interp(rho_L_node, InterpMode.MID)
    rho_g_node = params["rho_g_node"]
    rho_g_segment = grid_interp(rho_g_node, InterpMode.MID)
    

    # 体积分数相关变量
    alpha_L_node = params["alpha_L_node"]
    alpha_L_segment = grid_interp(alpha_L_node, InterpMode.MID)
    alpha_D_node = params["alpha_D_node"]
    alpha_D_segment = grid_interp(alpha_D_node, InterpMode.MID)
    alpha_g_node = params["alpha_g_node"]

    alpha_g_segment = grid_interp(alpha_g_node, InterpMode.MID)
    
    # 初始化变量
    rho_L_node_iteration_init = params["rho_L_node_iteration_init"]
    rho_L_segment_iteration_init = grid_interp(rho_L_node_iteration_init, InterpMode.MID)
    alpha_L_node_iteration_init = params["alpha_L_node_iteration_init"]

    alpha_L_segment_iteration_init = grid_interp(alpha_L_node_iteration_init, InterpMode.MID)
    V_L_segment_iteration_init = params["V_L_segment_iteration_init"]


    # 相间压力作用项，人工拟合项中的截面分数梯度项
    alpha_L_segment_diff = np.zeros_like(alpha_L_segment)
    alpha_L_segment_diff[1:-1] = difference(alpha_L_segment, dx, 2)
    alpha_L_segment_diff[0] = 2 * alpha_L_segment_diff[1] - alpha_L_segment_diff[2]
    alpha_L_segment_diff[-1] = 2 * alpha_L_segment_diff[-2] - alpha_L_segment_diff[-3]
    
    # 动量方程的源
    moment_src_L = (
            # 摩擦力
            - liquid_wall_friction_force / cross_section_area_segment
            + gas_liquid_friction_force / cross_section_area_segment
            # 相变
            - liquid_vapour_mass_transfer_rate * (alpha_L_segment / (alpha_L_segment + alpha_D_segment) * transition_velocity_from_liquid_film)
            - liquid_to_droplet_mass_transfer_rate * V_i_segment
            + droplet_to_liquid_mass_transfer_rate * droplet_velocity
            # 重力
            - (rho_L_segment * alpha_L_segment) * Physics.g.value * np.sin(tilt_angle_segment)
            # 气液界面压力
            # - alpha_L_segment * D_segment *(rho_L_segment - rho_g_segment) * Physics.g.value*alpha_L_segment_diff * np.cos(tilt_angle_segment)
            - 1.2
            * (alpha_g_segment * alpha_L_segment * rho_g_segment * rho_L_segment)
            / (alpha_g_segment * rho_L_segment + alpha_L_segment * rho_g_segment)
            * (V_g_segment - V_L_segment) ** 2
            * alpha_L_segment_diff
    )
    
    # 迎风格式
    alpha_rho_V_segment = alpha_L_segment * rho_L_segment * V_L_segment
    alpha_rho_V_L_plus = np.maximum(alpha_rho_V_segment, 0)
    alpha_rho_V_L_neg = -np.maximum(-alpha_rho_V_segment, 0)
    
    # 矩阵系数
    a_ww = (1/2*alpha_rho_V_L_plus)[1:-1]
    a_w = (-2*alpha_rho_V_L_plus)[1:-1]
    a_e = (3/2*alpha_rho_V_L_plus - 3/2*alpha_rho_V_L_neg + dx/dt*alpha_L_segment*rho_L_segment)[1:-1] 
    a_ee = (2*alpha_rho_V_L_neg)[1:-1]
    a_eee = (-1/2*alpha_rho_V_L_neg)[1:-1]
    b =(alpha_L_segment_iteration_init*rho_L_segment_iteration_init*V_L_segment_iteration_init*dx/dt + moment_src_L*dx - alpha_L_segment*difference(p_node, dx, 1)*dx)[1:-1]
    
    # 边界采用一阶迎风
    a_w[0] = - alpha_rho_V_L_plus[1]
    a_w[-1] = - alpha_rho_V_L_plus[-2]
    a_e[0] = alpha_rho_V_L_plus[1] - alpha_rho_V_L_neg[1] + dx/dt*alpha_L_segment[1]*rho_L_segment[1]
    a_e[-1] = alpha_rho_V_L_plus[-2] - alpha_rho_V_L_neg[-2] + dx/dt*alpha_L_segment[-2]*rho_L_segment[-2]
    a_ee[0] = alpha_rho_V_L_neg[1]
    a_ee[-1] = alpha_rho_V_L_neg[-2]
    a_ww[0] = 0
    a_ww[-1] = 0
    a_eee[0] = 0
    a_eee[-1] = 0
    
    # 添加边界的影响
    b[0] -= a_w[0] * V_L_segment[0]
    b[-1] -= a_ee[-1] * V_L_segment[-1]
    b[1] -= a_ww[1] * V_L_segment[0]
    b[-2] -= a_eee[-2] * V_L_segment[-1]
    
    # 求解
    V_L_segment_update = np.copy(V_L_segment)
    V_L_segment_update[1:-1] = penta_diagonal_solve(a_e, a_ee, a_w, a_eee, a_ww, b)
    return V_L_segment_update

def cal_moment_gD(params,
                  gas_wall_friction_force:np.ndarray, 
                  gas_liquid_friction_force:np.ndarray, 
                  liquid_vapour_mass_transfer_rate:np.ndarray, 
                  V_i_segment:np.ndarray, 
                  transition_velocity_from_droplet:np.ndarray,

                  liquid_to_droplet_mass_transfer_rate:np.ndarray, 
                  droplet_to_liquid_mass_transfer_rate:np.ndarray, 

                  droplet_velocity:np.ndarray) -> np.ndarray:
    """动量方程计算,气相和液滴相

    Args:
        params: dict 参数字典
        gas_wall_friction_force: np.ndarray 气相壁面摩擦力
        gas_liquid_friction_force: np.ndarray 气液摩擦力
        liquid_vapour_mass_transfer_rate: np.ndarray 液相蒸气质量传递速率
        V_i_segment: np.ndarray 液相到液滴质量传递速率
        transition_velocity_from_droplet: np.ndarray 来自液滴蒸发的相变速度
        liquid_to_droplet_mass_transfer_rate: np.ndarray 液相到液滴质量传递速率
        droplet_to_liquid_mass_transfer_rate: np.ndarray 液滴到液相质量传递速率

        droplet_velocity: np.ndarray 液滴速度

    Returns:
        V_g_segment_update: np.ndarray:, 气相速度
    """
    dx = params["dx"]
    dt = params["dt"]
    p_node = params["p_node"]
    cross_section_area_segment = params["cross_section_area_segment"]
    tilt_angle_segment = params["tilt_angle_segment"]
    


    # 速度相关变量
    V_g_segment = params["V_g_segment"]
    V_g_node = grid_extrap(V_g_segment, 1)
    V_L_segment = params["V_L_segment"]
    V_D_segment = params["V_D_segment"]
    V_D_node = grid_extrap(V_D_segment, 1)
    

    # 密度相关变量
    rho_g_node = params["rho_g_node"]
    rho_g_segment = grid_interp(rho_g_node, InterpMode.MID)
    rho_L_node = params["rho_L_node"]
    rho_L_segment = grid_interp(rho_L_node, InterpMode.MID)
    

    # 体积分数相关变量
    alpha_g_node = params["alpha_g_node"]
    alpha_g_segment = grid_interp(alpha_g_node, InterpMode.MID)
    alpha_L_node = params["alpha_L_node"]
    alpha_L_segment = grid_interp(alpha_L_node, InterpMode.MID)
    alpha_D_node = params["alpha_D_node"]

    alpha_D_segment = grid_interp(alpha_D_node, InterpMode.MID)
    
    # 初始化变量
    rho_g_node_iteration_init = params["rho_g_node_iteration_init"]
    rho_g_segment_iteration_init = grid_interp(rho_g_node_iteration_init, InterpMode.MID)
    rho_L_node_iteration_init = params["rho_L_node_iteration_init"]

    rho_L_segment_iteration_init = grid_interp(rho_L_node_iteration_init, InterpMode.MID)
    alpha_g_node_iteration_init = params["alpha_g_node_iteration_init"]
    alpha_g_segment_iteration_init = grid_interp(alpha_g_node_iteration_init, InterpMode.MID)

    alpha_D_node_iteration_init = params["alpha_D_node_iteration_init"]
    alpha_D_segment_iteration_init = grid_interp(alpha_D_node_iteration_init, InterpMode.MID)
    V_g_segment_iteration_init = params["V_g_segment_iteration_init"]
    V_D_segment_iteration_init = params["V_D_segment_iteration_init"]



    # 相间压力作用项，人工拟合项
    alpha_g_segment_diff = np.zeros_like(alpha_g_segment)
    alpha_g_segment_diff[1:-1] = difference(alpha_g_segment, dx, 2)
    alpha_g_segment_diff[0] = 2 * alpha_g_segment_diff[1] - alpha_g_segment_diff[2]
    alpha_g_segment_diff[-1] = 2 * alpha_g_segment_diff[-2] - alpha_g_segment_diff[-3]
    
    # 动量方程的源
    moment_src_gD = (
            # 摩擦力
            - gas_wall_friction_force / cross_section_area_segment
            - gas_liquid_friction_force / cross_section_area_segment
            # 相变
            + liquid_vapour_mass_transfer_rate * (alpha_L_segment / (alpha_L_segment + alpha_D_segment) * transition_velocity_from_droplet)
            + liquid_to_droplet_mass_transfer_rate * V_i_segment
            - droplet_to_liquid_mass_transfer_rate * droplet_velocity
            # 重力

            - (rho_g_segment * alpha_g_segment + alpha_D_segment * rho_L_segment)
            * Physics.g.value
            * np.sin(tilt_angle_segment)
            
            - 1.2
            * (alpha_g_segment * alpha_L_segment * rho_g_segment * rho_L_segment)
            / (alpha_g_segment * rho_L_segment + alpha_L_segment * rho_g_segment)
            * (V_g_segment - V_L_segment) ** 2
            * alpha_g_segment_diff
    )
    
    # 迎风格式
    alpha_rho_V_segment = alpha_g_segment * rho_g_segment * V_g_segment
    alpha_rho_V_g_plus = np.maximum(alpha_rho_V_segment, 0)
    alpha_rho_V_g_neg = -np.maximum(-alpha_rho_V_segment, 0)
    
    # 矩阵系数
    a_ww = (1/2*alpha_rho_V_g_plus)[1:-1]
    a_w = (-2*alpha_rho_V_g_plus)[1:-1]
    a_e = (3/2*alpha_rho_V_g_plus - 3/2*alpha_rho_V_g_neg + dx/dt*(alpha_g_segment*rho_g_segment + alpha_D_segment*rho_L_segment))[1:-1]
    a_ee = (2*alpha_rho_V_g_neg)[1:-1]
    a_eee = (-1/2*alpha_rho_V_g_neg)[1:-1]
    
    b = ((rho_g_segment_iteration_init*alpha_g_segment_iteration_init*V_g_segment_iteration_init + 
          rho_L_segment_iteration_init*alpha_D_segment_iteration_init*V_D_segment_iteration_init)*dx/dt + 
         moment_src_gD*dx - (alpha_g_segment + alpha_D_segment)*difference(p_node, dx, 1)*dx)[1:-1]
    
    # 边界采用一阶迎风
    a_w[0] = - alpha_rho_V_g_plus[1]
    a_w[-1] = - alpha_rho_V_g_plus[-2]
    a_e[0] = alpha_rho_V_g_plus[1] - alpha_rho_V_g_neg[1] + dx/dt*(alpha_g_segment[1]*rho_g_segment[1] + alpha_D_segment[1]*rho_L_segment[1])
    a_e[-1] = alpha_rho_V_g_plus[-2] - alpha_rho_V_g_neg[-2] + dx/dt*(alpha_g_segment[-2]*rho_g_segment[-2] + alpha_D_segment[-2]*rho_L_segment[-2])
    a_ee[0] = alpha_rho_V_g_neg[1]
    a_ee[-1] = alpha_rho_V_g_neg[-2]
    a_ww[0] = 0
    a_ww[-1] = 0
    a_eee[0] = 0
    a_eee[-1] = 0
    
    # 添加边界的影响
    b[0] -= a_w[0] * V_g_segment[0]
    b[-1] -= a_ee[-1] * V_g_segment[-1]
    b[1] -= a_ww[1] * V_g_segment[0]
    b[-2] -= a_eee[-2] * V_g_segment[-1]
    
    # 求解
    V_g_segment_update = np.copy(V_g_segment)
    V_g_segment_update[1:-1] = penta_diagonal_solve(a_e, a_ee, a_w, a_eee, a_ww, b)
    return V_g_segment_update
